'''
Created on Jul 2, 2011

@author: Nam Khanh
'''

from dep import *
from liblinear import *
from liblinearutil import *

n_gram = 3

def to_node(line):
    """
    """
    token = line.split('\t')
    idx = int(token[0])
    form = token[1]
    pos = token[4]
    
    #if training file
    if len(token) < 9:
        head_idx = int(token[6])
        return DepNode(idx = idx, form = form, pos = pos, head_idx = head_idx)
    else:
        target_head = int(token[8])
        return DepNode(idx = idx, form = form, pos = pos, target_head = target_head)

def get_feature(tree, word, history, symbol_table):
    """
    """
    feature = list()
    for i in range(0,n_gram):
        if word-i >= 0:
            f = 'wf_w-' + str(i) + '=' + tree.get_form(word-i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
            else:
                symbol_table[f] = len(symbol_table) + 1
                feature.append(symbol_table[f])
                
            f = 'pos_w-' + str(i) + '=' + tree.get_pos(word-i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
            else:
                symbol_table[f] = len(symbol_table) + 1
                feature.append(symbol_table[f])
                
        if history-i >= 0:
            f = 'wf_h-' + str(i) + '=' + tree.get_form(history-i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
            else:
                symbol_table[f] = len(symbol_table) + 1
                feature.append(symbol_table[f])
                
            f = 'pos_h-' + str(i) + '=' + tree.get_pos(history-i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
            else:
                symbol_table[f] = len(symbol_table) + 1
                feature.append(symbol_table[f])
            
    for i in range(0,n_gram):
        if word+i < len(tree):
            f = 'wf_w+' + str(i) + '=' + tree.get_form(word+i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
            else:
                symbol_table[f] = len(symbol_table) + 1
                feature.append(symbol_table[f])
                
            f = 'pos_w+' + str(i) + '=' + tree.get_pos(word+i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
            else:
                symbol_table[f] = len(symbol_table) + 1
                feature.append(symbol_table[f])
                
        if history+i < len(tree):
            f = 'wf_h+' + str(i) + '=' + tree.get_form(history+i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
            else:
                symbol_table[f] = len(symbol_table) + 1
                feature.append(symbol_table[f])
                
            f = 'pos_h+' + str(i) + '=' +  tree.get_pos(history+i)
            if symbol_table.has_key(f):
                feature.append(symbol_table[f])
            else:
                symbol_table[f] = len(symbol_table) + 1
                feature.append(symbol_table[f])
       
    f_dict = dict()
    for f in feature:
        f_dict[f] = 1
    
    return f_dict

def process(sentence, symbol_table, y, x):
    """
    """
    tree = DepTree([])
    root_node = DepNode(0, '#$ROOT$#', '#$ROOT$#', -1)
    tree.add(root_node)
    
    for line in sentence:
        tree.add(to_node(line))

    head = list()
    for idx in range(0, len(tree)):
        head.append(-1)

    for word in range(1, len(tree)):
        for i in range(1, word+1):
            history = word - i
            if (head[word] != -1 and head[history] != -1):
                continue
            trans = 0
            if tree.get_head_idx(word) == history:
                trans = 1
                head[word] = history
                
            if tree.get_head_idx(history) == word:
                trans = 2
                head[history] = word
            
            f = get_feature(tree, word, history, symbol_table)
#            print trans, f
            x.append(f)
            y.append(trans)            
                
def do_train(filename):
    """
    """
    symbol_table = dict()
    y = list()
    x = list()
    print "Selecting features n_gram = ", n_gram
    sentence = []
    number_sentence = 0
    with open(filename, 'r') as fin:
        for line in fin.readlines():
            line = line.strip()
            if (len(line) == 0):
                number_sentence = number_sentence + 1
                print "Processing sentence ", number_sentence
                process(sentence, symbol_table, y, x)
                sentence = []
            else:
                sentence.append(line)
    if len(sentence) > 0:
        number_sentence = number_sentence + 1
        print "Processing sentence ", number_sentence        
        process(sentence, symbol_table, y, x)

    with open('index_' + str(n_gram) + '.table', 'w') as fout:
        for k, v in symbol_table.items():
            fout.write(k)
            fout.write('\t')
            fout.write(str(v))
            fout.write('\n')

#    print y
#    print x            
    print "Training..."        
    prob = problem(y, x)
    param = parameter('-c 4')
    model = liblinear.train(prob, param)
    save_model('english_' + str(n_gram) + '.model', model)
    
if __name__=='__main__':
    if len(sys.argv) < 3:
        print '\nUsage: python {0} filename -n_gram'.format(sys.argv[0])
        sys.exit(1)
        
    n_gram = int(sys.argv[2])   
    do_train(sys.argv[1])